import sys
import time
import joblib
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from PySide6.QtCore import Qt
from PySide6.QtGui import QIcon, QPixmap
from PySide6.QtWidgets import QApplication, QWidget, QFileDialog, QTableWidgetItem, QMessageBox, QDialog, QVBoxLayout, \
    QLabel
from qfluentwidgets import InfoBar, InfoBarPosition
from sklearn.neighbors import KNeighborsClassifier
from View.ui_mainwindow import Ui_Form
from Resource import res_rc


class MainWindow(QWidget):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.ui = Ui_Form()
        self.ui.setupUi(self)
        self.setup_connections()
        self.knn_model = None  # 初始化模型变量
        # 数据导入状态
        self.train_data_loaded = False
        self.test_data_loaded = False

    def setup_connections(self):  # 连接信号与槽
        self.ui.btn_loadTrain_data.clicked.connect(lambda: self.load_data("train"))
        self.ui.btn_loadTest_data.clicked.connect(lambda: self.load_data("test"))
        self.ui.btn_train.clicked.connect(self.train_knn_model)
        self.ui.btn_indentify.clicked.connect(self.on_btn_identify_clicked)
        self.ui.btn_ana.clicked.connect(self.plot_analysis)

    def load_data(self, data_type):  # 数据导入函数
        if data_type == "test" and not self.knn_model:
            self.createErrorInfoBar("请在训练模型后再导入测试数据！")
            return
        file_path, _ = QFileDialog.getOpenFileName(self, "选择CSV文件", "", "CSV Files (*.csv);;All Files (*)")
        if file_path:
            df = pd.read_csv(file_path, header=None)
            df[4] = df[4].astype(str)
            self.ui.dataTable.setRowCount(0)
            self.ui.dataTable.setColumnCount(len(df.columns))
            for i, row in df.iterrows():
                self.ui.dataTable.insertRow(i)
                for j, value in enumerate(row):
                    item = QTableWidgetItem(str(value))
                    item.setTextAlignment(Qt.AlignCenter)
                    self.ui.dataTable.setItem(i, j, item)
            self.ui.dataTable.resizeColumnsToContents()
            # 根据导入的数据类型设置导入状态
            if data_type == "train":
                self.train_data_loaded = True
            elif data_type == "test":
                self.test_data_loaded = True

    def check_data_loaded(self, data_type):  # 检查数据是否加载
        if data_type == "train" and not self.train_data_loaded:
            self.createErrorInfoBar("训练模型前请先加载训练数据！")
            return False
        elif data_type == "test" and not self.test_data_loaded:
            self.createErrorInfoBar("还未加载测试数据！")
            return False
        return True

    def extract_data(self):  # 这里从Table中提取数据
        row_count = self.ui.dataTable.rowCount()
        col_count = self.ui.dataTable.columnCount()
        data = []
        labels = []
        for i in range(row_count):
            row_data = [float(self.ui.dataTable.item(i, j).text()) for j in range(col_count - 1)]
            data.append(row_data)
            labels.append(int(self.ui.dataTable.item(i, col_count - 1).text()))
        return pd.DataFrame(data), pd.Series(labels)

    def train_knn_model(self):
        # 检查训练数据是否导入
        if not self.check_data_loaded("train"):
            return
        # 训练KNN模型
        data, labels = self.extract_data()
        k = self.ui.k_spinbox.value()
        if k < 1:
            QMessageBox.warning(self, "警告", "K值必须大于0")
            return
        self.knn_model = KNeighborsClassifier(n_neighbors=k)
        self.knn_model.fit(data, labels)
        joblib.dump(self.knn_model, './Resource/knn_model.pkl')
        time.sleep(1)
        self.createSuccessInfoBar("模型训练成功！")

    def on_btn_identify_clicked(self):  # 识别数据
        # 检查测试数据是否导入
        if not self.check_data_loaded("test"):
            return
        data, true_labels = self.extract_data()
        predictions = self.knn_model.predict(data)
        correct_count = 0
        for i, (prediction, true_label) in enumerate(zip(predictions, true_labels)):
            icon_path = ':/icon/green_correct.png' if prediction == true_label else ':/icon/red_error.png'
            icon = QIcon(QPixmap(icon_path))
            icon_item = QTableWidgetItem(str(true_label))
            icon_item.setIcon(icon)
            icon_item.setTextAlignment(Qt.AlignCenter)
            text = f"Prediction:{prediction} | {true_label}"
            # icon_item.setData(Qt.DisplayRole, text)
            icon_item.setText(text)
            self.ui.dataTable.setItem(i, self.ui.dataTable.columnCount() - 1, icon_item)  # 在分类标签列设置图标
            correct_count += int(prediction == true_label)
        accuracy = correct_count / len(true_labels) * 100
        self.ui.acc_linEdit.setText(f"{accuracy:.2f}%")

    def plot_analysis(self):
        # 检查训练数据是否导入
        if not self.check_data_loaded("train"):
            return
        data, labels = self.extract_data()
        species = ['0-setosa', '1-versicolor', '2-virginica']
        plt.figure(figsize=(10, 8))
        # 定义颜色映射，将0映射为红色，1映射为蓝色，2映射为绿色
        colors = ['red', 'blue', 'green']
        # 将 labels 转换为 numpy 数组
        labels = np.array(labels)
        # 子图1: Sepal Length vs Sepal Width
        plt.subplot(2, 1, 1)
        for i, species_name in enumerate(species):
            plt.scatter(data.iloc[labels == i, 0], data.iloc[labels == i, 1],
                        color=colors[i], label=species_name)
        plt.title('Sepal Length vs Sepal Width')
        plt.xlabel('Sepal Length')
        plt.ylabel('Sepal Width')
        plt.legend(title='Species')
        # 子图2: Petal Length vs Petal Width
        plt.subplot(2, 1, 2)
        for i, species_name in enumerate(species):
            plt.scatter(data.iloc[labels == i, 2], data.iloc[labels == i, 3],
                        color=colors[i], label=species_name)
        plt.title('Petal Length vs Petal Width')
        plt.xlabel('Petal Length')
        plt.ylabel('Petal Width')
        plt.legend(title='Species')
        plt.tight_layout()
        plt.savefig('./Resource/analysis.png')
        self.show_image_dialog('./Resource/analysis.png')

    def show_image_dialog(self, image_path):
        dialog = QDialog(self)
        dialog.setWindowTitle('分析结果')
        layout = QVBoxLayout()
        label = QLabel()
        pixmap = QPixmap(image_path)
        label.setPixmap(pixmap)
        layout.addWidget(label)
        dialog.setLayout(layout)
        dialog.exec()

    def createSuccessInfoBar(self, content):
        InfoBar.success(
            title='success~ ',
            content=content,
            orient=Qt.Horizontal,
            isClosable=True,
            position=InfoBarPosition.TOP,
            duration=4000,
            parent=self
        )

    def createErrorInfoBar(self, content):
        InfoBar.error(
            title='error~ ',
            content=content,
            orient=Qt.Horizontal,
            isClosable=True,
            position=InfoBarPosition.TOP,
            duration=4000,
            parent=self
        )

if __name__ == '__main__':
    app = QApplication(sys.argv)
    win = MainWindow()
    win.show()
    sys.exit(app.exec())
